import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import math
#--------------生成数据------------------

# 翻译：生成1000个样本数据
SAMPLES = 1000

# 设置一个“种子”值，这样每次运行这个notebook时都能得到相同的随机数。任何数字都可以用作种子值。
SEED = 1337
np.random.seed(SEED) # 这个函数设置随机数种子
tf.random.set_seed(SEED)

# 翻译：生成一个均匀分布的随机数，范围从0到2π，覆盖了完整的正弦波振荡
# unifotm函数：从给定的上下限之间均匀分布随机抽取一个值
x_values = np.random.uniform(low=0, high=2*math.pi, size=SAMPLES)

# 打乱数据顺序
np.random.shuffle(x_values)

# sin函数：计算正弦值
y_values = np.sin(x_values)

# 给每个y值添加一个小的随机数
y_values += 0.1 * np.random.randn(*y_values.shape)

# 绘制数据点
# plt.title("Sine Wave Data")
# plt.plot(x_values, y_values, 'b.')
# plt.show()

#--------------划分数据集-----------------------
# We'll use 60% of our data for training and 20% for testing. The remaining 20%
# will be used for validation. Calculate the indices of each section.
# 翻译：将数据集分为训练集、验证集和测试集
# 训练集：60%
# 验证集：20%
# 测试集：20%
TRAIN_SPLIT =  int(0.6 * SAMPLES)  # 训练集的索引 600
TEST_SPLIT = int(0.2 * SAMPLES + TRAIN_SPLIT) # 测试集的索引 800
# split函数，参数：数组、分割位置，返回值：分割后的数组 分割点在600  800 分成三段 0-600 600-800  800-1000
x_train, x_validate, x_test = np.split(x_values, [TRAIN_SPLIT, TEST_SPLIT]) # 训练集、验证集、测试集的x值
y_train, y_validate, y_test = np.split(y_values, [TRAIN_SPLIT, TEST_SPLIT]) # 训练集、验证集、测试集的y值

# plt.plot(x_train, y_train, 'b.', label="Train")
# plt.plot(x_validate, y_validate, 'y.', label="Validate")
# plt.plot(x_test, y_test, 'r.', label="Test")
# plt.legend()
# plt.show()

print("end\n")


# --------------定义模型-------------
import tensorflow as tf

# 使用Keras创建一个简单的模型架构
model_1 = tf.keras.Sequential()

# 第一层接受一个标量输入，并通过16个“神经元”传递它。
# 输入的标准量是1，因为我们只有一个输入特征。
# 神经元根据'relu'激活函数决定是否激活。
model_1.add(tf.keras.layers.Dense(16, activation='relu', input_shape=(1,)))

# 第二层是一个具有16个神经元的全连接层，并使用'relu'激活函数。
model_1.add(tf.keras.layers.Dense(16, activation='relu'))

# 第三层是一个具有16个神经元的全连接层，并使用'relu'激活函数。
model_1.add(tf.keras.layers.Dense(16, activation='relu'))

# 最后一层是一个单一的神经元，因为我们希望输出一个单一的值
model_1.add(tf.keras.layers.Dense(1))

# 使用标准的优化器和损失函数编译模型进行回归
model_1.compile(optimizer='rmsprop', loss='mse', metrics=['mae'])

# 打印模型的架构摘要
model_1.summary()

# --------------训练模型-------------
# 训练模型，将训练集作为输入，将目标值作为输出
# 参数：训练集、目标值、轮数、批次大小、验证集
history_1 = model_1.fit(x_train, y_train, epochs=2000, batch_size=16,
                    validation_data=(x_validate, y_validate))

# 在0-2pi范围内生成1000个样本数据
x_test_values = np.random.uniform(low=0, high=2*math.pi, size=1000)

# 预测模型输出
y_test_predictions = model_1.predict(x_test_values)

# 绘制预测值和实际值
plt.clf()
plt.title('Comparison of predicted values against actual values')
plt.plot(x_test_values, y_test_predictions, 'bo', label='Predicted')
plt.plot(x_test_values, np.sin(x_test_values), 'r.', label='Actual')
plt.legend() # 显示图例
plt.show() # 显示图形


# 翻译：将模型转换为TensorFlow Lite格式，不进行量化
converter = tf.lite.TFLiteConverter.from_keras_model(model_1)
tflite_model = converter.convert()
# 翻译：将模型保存到磁盘
open("sine_model.tflite", "wb").write(tflite_model)


# 翻译：将模型转换为TensorFlow Lite格式，并进行8位的量化，这是量化成8位的过程，可以减少模型大小并提高性能。
# 保证输入跟输出都是int8，因为量化后的模型大小会减小。
converter = tf.lite.TFLiteConverter.from_keras_model(model_1)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
# 定义一个生成器来提供测试数据x的值，作为代表性数据集，并告诉转换器使用它
def representative_dataset_gen():
  for value in x_test:
    yield [np.array([[value]], dtype=np.float32)] # 这dtepe=np.float32是为了保证输入的数据类型为float32

converter.representative_dataset = representative_dataset_gen
tflite_model_quantized = converter.convert()
# 转换模型
tflite_model = converter.convert()
# 翻译：将量化后的模型保存到磁盘
open("sine_model_quantized_generated.tflite", "wb").write(tflite_model)




# Instantiate an interpreter for each model
sine_model = tf.lite.Interpreter('sine_model.tflite')
sine_model_quantized = tf.lite.Interpreter('sine_model_quantized_generated.tflite')

# Allocate memory for each model
sine_model.allocate_tensors()
sine_model_quantized.allocate_tensors()

# Get indexes of the input and output tensors
sine_model_input_index = sine_model.get_input_details()[0]["index"]
sine_model_output_index = sine_model.get_output_details()[0]["index"]
sine_model_quantized_input_index = sine_model_quantized.get_input_details()[0]["index"]
sine_model_quantized_output_index = sine_model_quantized.get_output_details()[0]["index"]

# Create arrays to store the results
sine_model_predictions = []
sine_model_quantized_predictions = []

# Run each model's interpreter for each value and store the results in arrays
for x_value in x_test:
  # Create a 2D tensor wrapping the current x value
  x_value_tensor = tf.convert_to_tensor([[x_value]], dtype=np.float32)
  # Write the value to the input tensor
  sine_model.set_tensor(sine_model_input_index, x_value_tensor)
  # Run inference
  sine_model.invoke()
  # Read the prediction from the output tensor
  sine_model_predictions.append(
      sine_model.get_tensor(sine_model_output_index)[0])
  # Do the same for the quantized model
  # sine_model_quantized.set_tensor(sine_model_quantized_input_index, x_value_tensor)
  sine_model_quantized.invoke()
  sine_model_quantized_predictions.append(
      sine_model_quantized.get_tensor(sine_model_quantized_output_index)[0])


# See how they line up with the data
plt.clf()
plt.title('Comparison of various models against actual values')
plt.plot(x_test, y_test, 'bo', label='Actual')
# plt.plot(x_test, predictions, 'ro', label='Original predictions')
plt.plot(x_test, sine_model_predictions, 'ro', label='Lite predictions')
plt.plot(x_test, sine_model_quantized_predictions, 'gx', label='Lite quantized predictions')
plt.legend()
plt.show()


